{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[View the runnable example on GitHub](https://github.com/intel-analytics/BigDL/tree/main/python/nano/tutorial/notebook/inference/pytorch/accelerate_pytorch_inference_async_pipeline.ipynb)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Accelerate PyTorch Inference using Async Multi-stage Pipeline"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If there are multiple stages (e.g., preprocessing, model inferencing, postprocessing, ...) in your AI application, an asynchronized pipeline that parallelize different stages could significantly improve the inference throughput of the whole pipeline.\n",
    "\n",
    "In a pipeline with N stages could work like this, where multiple (not necessarily to be the same number) workers could be enabled for each stage. The data will flow along each stages.\n",
    "\n",
    "<img width=\"925\" alt=\"image\" src=\"https://user-images.githubusercontent.com/35031544/236608726-072b1fbb-c920-427f-ae24-57565f2c4e72.png\">"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "nbsphinx": "hidden"
   },
   "source": [
    "To apply async multi-stage inference pipeline acceleration, you need to install our lately BigDL-Nano for PyTorch inference first:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbsphinx": "hidden"
   },
   "outputs": [],
   "source": [
    "!pip install --pre --upgrade bigdl-nano[pytorch,inference] # install the nightly-built version\n",
    "!source bigdl-nano-init"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "nbsphinx": "hidden"
   },
   "source": [
    "> 📝 **Note**\n",
    ">\n",
    "> We recommend to run the commands above, especially `source bigdl-nano-init` before jupyter kernel is started, or some of the optimizations may not take effect."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's take an [ResNet-18 model](https://pytorch.org/vision/main/models/generated/torchvision.models.resnet18.html) pretrained on ImageNet dataset as an example. First, we load the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.models import resnet18\n",
    "\n",
    "model = resnet18(pretrained=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we create our preprocessing function to transform image samples to tensors in our desired forms. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torchvision import transforms\n",
    "\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize(256),\n",
    "    transforms.ColorJitter(),\n",
    "    transforms.Resize(224),\n",
    "    transforms.ToTensor(),\n",
    "])\n",
    "\n",
    "\n",
    "def preprocess(ori_img):\n",
    "    img = transform(ori_img)\n",
    "    batch = torch.stack([img for _ in range(8)], 0)\n",
    "    return batch"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To enable async multi-stage inference pipeline acceleration for your PyTorch model, **the major work you need to make is to import BigDL-Nano** `Pipeline`**, and add each stage by specifying the stage name, the corresponding function and configs**:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bigdl.nano.pytorch.inference.pipeline import Pipeline\n",
    "\n",
    "pipeline = Pipeline([\n",
    "    (\"preprocess\", preprocess,  {\"cores_per_worker\": 4, 'worker_num': 4}),\n",
    "    (\"inference\", model,  {\"cores_per_worker\": 4, 'worker_num': 4}),\n",
    "])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> 📝 **Note**\n",
    "> \n",
    "> For configs, `cores_per_worker` specifies the number of physical cores for each worker, default set to largest core numbers of your system. And `worker_num` specifies the number of workers, default set to 1. "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then you can feed your image samples to `pipeline.run()` API to start multi-stage pipeline inferencing. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from PIL import Image\n",
    "\n",
    "imarray = np.random.rand(100, 100, 3) * 255\n",
    "test_image = Image.fromarray(imarray.astype('uint8')).convert('RGB')\n",
    "\n",
    "input_samples = [test_image for _ in range(100)]\n",
    "pipeline.run(input_samples)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> 📚 **Related Readings**\n",
    "> \n",
    "> - [How to install BigDL-Nano](https://bigdl.readthedocs.io/en/latest/doc/Nano/Overview/install.html)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nano-pytorch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.7.15 (default, Nov 24 2022, 21:12:53) \n[GCC 11.2.0]"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "8a5edab282632443219e051e4ade2d1d5bbc671c781051bf1437897cbdfea0f1"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
